// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

package com.microsoft.azure.keyvault.cryptography;

import java.io.IOException;
import java.security.GeneralSecurityException;
import java.security.InvalidAlgorithmParameterException;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.Provider;
import java.security.Security;
import java.security.interfaces.ECPublicKey;
import java.security.spec.ECGenParameterSpec;
import java.security.spec.ECParameterSpec;
import java.security.spec.EllipticCurve;
import java.security.spec.InvalidKeySpecException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.UUID;

import org.apache.commons.lang3.tuple.Pair;
import org.apache.commons.lang3.tuple.Triple;

import com.google.common.collect.ImmutableMap;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.microsoft.azure.keyvault.core.IKey;
import com.microsoft.azure.keyvault.cryptography.algorithms.Ecdsa;
import com.microsoft.azure.keyvault.cryptography.algorithms.Es256k;
import com.microsoft.azure.keyvault.cryptography.algorithms.Es256;
import com.microsoft.azure.keyvault.cryptography.algorithms.Es384;
import com.microsoft.azure.keyvault.cryptography.algorithms.Es512;
import com.microsoft.azure.keyvault.webkey.JsonWebKey;
import com.microsoft.azure.keyvault.webkey.JsonWebKeyCurveName;


public class EcKey implements IKey {

    public static final String P256 = "secp256r1";
    public static final String P384 = "secp384r1";
    public static final String P521 = "secp521r1";
    public static final String P256K = "secp256k1";
    public static final Map<JsonWebKeyCurveName, String> CURVE_TO_SIGNATURE = ImmutableMap.<JsonWebKeyCurveName, String>builder()
            .put(JsonWebKeyCurveName.P_256, Es256.ALGORITHM_NAME)
            .put(JsonWebKeyCurveName.P_384, Es384.ALGORITHM_NAME)
            .put(JsonWebKeyCurveName.P_521, Es512.ALGORITHM_NAME)
            .put(JsonWebKeyCurveName.P_256K, Es256k.ALGORITHM_NAME)
            .build();
    public static final Map<JsonWebKeyCurveName, String> CURVE_TO_SPEC_NAME = ImmutableMap.<JsonWebKeyCurveName, String>builder()
            .put(JsonWebKeyCurveName.P_256, P256)
            .put(JsonWebKeyCurveName.P_384, P384)
            .put(JsonWebKeyCurveName.P_521, P521)
            .put(JsonWebKeyCurveName.P_256K, P256K)
            .build();


    private final String kid;
    private final KeyPair keyPair;
    private final Provider provider;
    private final JsonWebKeyCurveName curve;

    protected final String signatureAlgorithm;
    protected String defaultEncryptionAlgorithm;

    public static JsonWebKeyCurveName getDefaultCurve() {
        return JsonWebKeyCurveName.P_256;
    }

    /**
     * Constructor.
     *
     * Generates a new EcKey with a P_256 curve and a randomly generated kid.
     *
     * @throws NoSuchAlgorithmException
     * @throws InvalidAlgorithmParameterException
     */
    public EcKey() throws NoSuchAlgorithmException, InvalidAlgorithmParameterException {
        this(UUID.randomUUID().toString());
    }

    /**
     * Constructor.
     *
     * Generates a new EcKey with a P_256 curve and the given kid.
     *
     * @param kid
     * @throws NoSuchAlgorithmException
     * @throws InvalidAlgorithmParameterException
     */
    public EcKey(String kid) throws NoSuchAlgorithmException, InvalidAlgorithmParameterException {
        this(kid, getDefaultCurve(), Security.getProvider("SunEC"));
    }

    /**
     * Constructor.
     *
     * Generates a new EcKey with the given curve and kid.
     * @param kid
     * @param curve
     * @throws NoSuchAlgorithmException
     * @throws InvalidAlgorithmParameterException
     */
    public EcKey(String kid, JsonWebKeyCurveName curve) throws NoSuchAlgorithmException, InvalidAlgorithmParameterException {
        this(kid, curve, Security.getProvider("SunEC"));
    }

    /**
     * Constructor.
     *
     * Generates a new EcKey with the given curve and kid.
     * @param kid
     * @param curve
     * @param provider Java security provider
     * @throws InvalidAlgorithmParameterException
     * @throws NoSuchAlgorithmException
     */
    public EcKey(String kid, JsonWebKeyCurveName curve, Provider provider) throws InvalidAlgorithmParameterException, NoSuchAlgorithmException {
        this.kid = kid;
        this.provider = provider;
        this.curve = curve;

        signatureAlgorithm = CURVE_TO_SIGNATURE.get(curve);
        if (signatureAlgorithm == null) {
            throw new NoSuchAlgorithmException("Curve not supported.");
        }

        final KeyPairGenerator generator = KeyPairGenerator.getInstance("EC", provider);
        ECGenParameterSpec gps = new ECGenParameterSpec(CURVE_TO_SPEC_NAME.get(curve));

        generator.initialize(gps);
        keyPair = generator.generateKeyPair();

    }

    /**
     * Constructor.
     *
     * Generates a new EcKey with the given keyPair.
     * The keyPair must be an ECKey.
     * @param kid
     * @param keyPair
     * @throws NoSuchAlgorithmException
     * @throws InvalidAlgorithmParameterException
     */
    public EcKey(String kid, KeyPair keyPair) throws NoSuchAlgorithmException, InvalidAlgorithmParameterException {
        this(kid, keyPair, Security.getProvider("SunEC"));
    }

    /**
     * Constructor.
     *
     * Generates a new EcKey with the given keyPair.
     * The keyPair must be an ECKey.
     * @param kid
     * @param keyPair
     * @param provider Java security provider
     * @throws NoSuchAlgorithmException
     * @throws InvalidAlgorithmParameterException
     */
    public EcKey(String kid, KeyPair keyPair, Provider provider) throws NoSuchAlgorithmException, InvalidAlgorithmParameterException {

        if (Strings.isNullOrWhiteSpace(kid)) {
            throw new IllegalArgumentException("Please provide a kid");
        }

        if (keyPair == null) {
            throw new IllegalArgumentException("Please provide an ECKey");
        }

        if (keyPair.getPublic() == null || !(keyPair.getPublic() instanceof ECPublicKey)) {
            throw new IllegalArgumentException("The keyPair provided is not an ECKey");
        }

        this.kid = kid;
        this.keyPair = keyPair;
        this.provider = provider;
        curve = getCurveFromKeyPair(keyPair);
        signatureAlgorithm = CURVE_TO_SIGNATURE.get(curve);
        if (signatureAlgorithm == null) {
            throw new IllegalArgumentException("Curve not supported.");
        }
    }

    /**
     * Converts JSON web key to EC key pair, does not include the private key.
     * @param jwk
     * @return EcKey
     * @throws NoSuchAlgorithmException
     * @throws InvalidAlgorithmParameterException
     * @throws InvalidKeySpecException
     * @throws NoSuchProviderException
     */
    public static EcKey fromJsonWebKey(JsonWebKey jwk) throws NoSuchAlgorithmException, InvalidAlgorithmParameterException, InvalidKeySpecException, NoSuchProviderException {
        return fromJsonWebKey(jwk, false, null);
    }

    /**
     * Converts JSON web key to EC key pair and include the private key if set to true.
     * @param jwk
     * @param includePrivateParameters true if the EC key pair should include the private key. False otherwise.
     * @return EcKey
     * @throws NoSuchAlgorithmException
     * @throws InvalidAlgorithmParameterException
     * @throws InvalidKeySpecException
     * @throws NoSuchProviderException
     */
    public static EcKey fromJsonWebKey(JsonWebKey jwk, boolean includePrivateParameters) throws NoSuchAlgorithmException, InvalidAlgorithmParameterException, InvalidKeySpecException, NoSuchProviderException {
        return fromJsonWebKey(jwk, includePrivateParameters, null);
    }

    /**
     * Converts JSON web key to EC key pair and include the private key if set to true.
     * @param jwk
     * @param includePrivateParameters true if the EC key pair should include the private key. False otherwise.
     * @param provider the Java Security Provider
     * @return EcKey
     */
    public static EcKey fromJsonWebKey(JsonWebKey jwk, boolean includePrivateParameters, Provider provider) {
        try {
            if (jwk.kid() != null) {
                return new EcKey(jwk.kid(), jwk.toEC(includePrivateParameters, provider));
            } else {
                throw new IllegalArgumentException("Json Web Key should have a kid");
            }
        } catch (GeneralSecurityException e) {
            throw new IllegalStateException(e);
        }
    }

    /**
     * Converts EcKey to JSON web key.
     * @return
     */
    public JsonWebKey toJsonWebKey() {
        return JsonWebKey.fromEC(keyPair, provider);
    }

    // Matches the curve of the keyPair to supported curves.
    private JsonWebKeyCurveName getCurveFromKeyPair(KeyPair keyPair) {
        try {
            ECPublicKey key = (ECPublicKey) keyPair.getPublic();
            ECParameterSpec spec = key.getParams();
            EllipticCurve crv = spec.getCurve();

            List<JsonWebKeyCurveName> curveList = Arrays.asList(JsonWebKeyCurveName.P_256, JsonWebKeyCurveName.P_384, JsonWebKeyCurveName.P_521, JsonWebKeyCurveName.P_256K);

            for (JsonWebKeyCurveName curve : curveList) {
                ECGenParameterSpec gps = new ECGenParameterSpec(CURVE_TO_SPEC_NAME.get(curve));
                KeyPairGenerator kpg = KeyPairGenerator.getInstance("EC", provider);
                kpg.initialize(gps);

                // Generate dummy keypair to get parameter spec.
                KeyPair apair = kpg.generateKeyPair();
                ECPublicKey apub = (ECPublicKey) apair.getPublic();
                ECParameterSpec aspec = apub.getParams();
                EllipticCurve acurve = aspec.getCurve();

                //Matches the parameter spec
                if (acurve.equals(crv)) {
                    return curve;
                }
            }

            //Did not find a supported curve.
            throw new IllegalArgumentException("Curve not supported.");
        } catch (GeneralSecurityException e) {
            throw new IllegalStateException(e);
        }
    }

    /**
     * @return curve of the key
     */
    public JsonWebKeyCurveName getCurve() {
        return curve;
    }

    /**
     *
     * @return the underlying keyPair of the key
     */
    public KeyPair getKeyPair() {
        return keyPair;
    }

    @Override
    public void close() throws IOException {
        // Intentionally empty
    }

    @Override
    public String getDefaultEncryptionAlgorithm() {
        return null;
    }

    @Override
    public String getDefaultKeyWrapAlgorithm() {
        return null;
    }

    @Override
    public String getDefaultSignatureAlgorithm() {
        return signatureAlgorithm;
    }

    @Override
    public String getKid() {
        return kid;
    }

    @Override
    public ListenableFuture<byte[]> decryptAsync(byte[] ciphertext, byte[] iv, byte[] authenticationData,
            byte[] authenticationTag, String algorithm) throws NoSuchAlgorithmException {
        throw new UnsupportedOperationException("Decrypt Async is not supported");
    }

    @Override
    public ListenableFuture<Triple<byte[], byte[], String>> encryptAsync(byte[] plaintext, byte[] iv,
            byte[] authenticationData, String algorithm) throws NoSuchAlgorithmException {
        throw new UnsupportedOperationException("Encrypt Async is not supported");
    }

    @Override
    public ListenableFuture<Pair<byte[], String>> wrapKeyAsync(byte[] key, String algorithm)
            throws NoSuchAlgorithmException {
        throw new UnsupportedOperationException("Wrap key is not supported");
    }

    @Override
    public ListenableFuture<byte[]> unwrapKeyAsync(byte[] encryptedKey, String algorithm)
            throws NoSuchAlgorithmException {
        throw new UnsupportedOperationException("Unwrap key is not supported");
    }

    @Override
    public ListenableFuture<Pair<byte[], String>> signAsync(byte[] digest, String algorithm) throws NoSuchAlgorithmException {
        
        if (keyPair.getPrivate() == null) {
            throw new UnsupportedOperationException("Sign is not supported without a private key.");
        }

        if (digest == null) {
            throw new IllegalArgumentException("Please provide a digest to sign.");
        }

        if (Strings.isNullOrWhiteSpace(algorithm)) {
            throw new IllegalArgumentException("Please provide a signature algorithm to use.");
        }

        // Interpret the requested algorithm
        Algorithm baseAlgorithm = AlgorithmResolver.Default.get(algorithm);

        if (baseAlgorithm == null || !(baseAlgorithm instanceof AsymmetricSignatureAlgorithm)) {
            throw new NoSuchAlgorithmException(algorithm);
        }
       
        Ecdsa algo = (Ecdsa) baseAlgorithm;
        ISignatureTransform signer = algo.createSignatureTransform(keyPair, provider);

        try {
            return Futures.immediateFuture(Pair.of(signer.sign(digest), algorithm));
        } catch (Exception e) {
            return Futures.immediateFailedFuture(e);
        }
    }

    @Override
    public ListenableFuture<Boolean> verifyAsync(byte[] digest, byte[] signature, String algorithm) throws NoSuchAlgorithmException {

        if (digest == null) {
            throw new IllegalArgumentException("Please provide a digest input.");
        }

        if (Strings.isNullOrWhiteSpace(algorithm)) {
            throw new IllegalArgumentException("Please provide an algorithm");
        }

        // Interpret the requested algorithm
        Algorithm baseAlgorithm = AlgorithmResolver.Default.get(algorithm);
        
        if (baseAlgorithm == null || !(baseAlgorithm instanceof AsymmetricSignatureAlgorithm)) {
            throw new NoSuchAlgorithmException(algorithm);
        }
        
        Ecdsa algo = (Ecdsa) baseAlgorithm;

        ISignatureTransform signer = algo.createSignatureTransform(keyPair, provider);
        
        try {
            return Futures.immediateFuture(signer.verify(digest, signature));
        } catch (Exception e) {
            return Futures.immediateFailedFuture(e);
        }
    }

}
